#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
IR 基础语句节点定义
"""

import ast
from dataclasses import dataclass, field
from typing import List, Iterator, Optional
from ..base import IRNode
from ...symbol_table import Symbol, SymbolKind


def build_stmt_from_ast(node: ast.AST, parent: Optional["IRNode"] = None) -> "IRNode":
    """从 AST 节点构建语句 IR 节点的统一函数

    Args:
        node: AST 节点
        parent: 父节点，如果提供则自动设置为构建节点的父节点

    Returns:
        对应的 IR 语句节点
    """
    if isinstance(node, ast.Return):
        from .control import IRReturn

        stmt = IRReturn(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Break):
        from .control import IRBreak

        stmt = IRBreak(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Continue):
        from .control import IRContinue

        stmt = IRContinue(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Assign):
        from .assignment import IRAssign

        stmt = IRAssign(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.AnnAssign):
        from .assignment import IRAssign

        stmt = IRAssign(parent=parent)
        stmt.build_from_ann_assign_ast(node)
    elif isinstance(node, ast.AugAssign):
        from .assignment import IRAugAssign

        stmt = IRAugAssign(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.If):
        from .control import IRIf

        stmt = IRIf(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Global):
        from .global_stmt import IRGlobal

        stmt = IRGlobal(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Nonlocal):
        from .global_stmt import IRNonlocal

        stmt = IRNonlocal(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.For):
        # 根据迭代器类型选择 IRForRange 或 IRForEach
        from .control import IRForRange, IRForEach

        if (
            isinstance(node.iter, ast.Call)
            and isinstance(node.iter.func, ast.Name)
            and node.iter.func.id == "range"
        ):
            stmt = IRForRange(parent=parent)
            if not stmt.build_from_ast(node):
                # 如果 IRForRange 构建失败，回退到 IRForEach
                stmt = IRForEach(parent=parent)
                stmt.build_from_ast(node)
        else:
            stmt = IRForEach(parent=parent)
            stmt.build_from_ast(node)
    elif isinstance(node, ast.While):
        from .control import IRWhile

        stmt = IRWhile(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Pass):
        # Pass 语句转换为专门的 IRPass 节点
        from .pass_stmt import IRPass

        stmt = IRPass(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Raise):
        from .exception import IRRaise

        stmt = IRRaise(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
        # 嵌套函数定义
        from ..declarations import IRFunctionNode

        stmt = IRFunctionNode(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Expr):
        from .expression import IRExprStmt

        stmt = IRExprStmt(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Delete):
        from .delete import IRDelete

        stmt = IRDelete(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, ast.Try):
        from .exception import IRTryExcept

        stmt = IRTryExcept(parent=parent)
        stmt.build_from_ast(node)
    elif isinstance(node, (ast.Import, ast.ImportFrom)):
        from ..declarations import IRImportNode

        expr = IRImportNode(parent=parent)
        expr.build_from_ast(node)
    else:
        # 对于不支持的节点类型，抛出异常
        raise ValueError(f"Unsupported AST node type: {type(node).__name__}")

    return stmt


@dataclass
class IRStmt(IRNode):
    """语句基类"""

    node_type: str = "stmt"


@dataclass
class IRStmtBlock(IRNode):
    """语句块类，用于统一处理所有多条语句场景（函数体、if语句块、循环体等）"""

    statements: List[IRNode] = field(default_factory=list)  # 存储 IR 语句节点
    node_type: str = "stmt_block"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        for stmt in self.statements:
            yield stmt
        yield from super().get_child_nodes()

    def build_from_ast(self, body: List[ast.stmt]) -> None:
        """从 AST 语句列表构建语句块内容

        Args:
            body: AST 语句列表
        """
        self.statements = self._extract_body_statements(body)

    def add_statement(self, stmt: IRNode) -> None:
        """添加语句到语句块"""
        self.statements.append(stmt)

    def is_empty(self) -> bool:
        """检查语句块是否为空"""
        return len(self.statements) == 0

    def get_statement_count(self) -> int:
        """获取语句数量"""
        return len(self.statements)

    def get_local_symbols(self) -> Iterator[Symbol]:
        """获取语句块中定义的所有本地符号

        遍历语句块中的所有赋值语句，收集局部变量定义
        同时递归遍历嵌套的语句块（如 if 语句的分支）

        注意：这个方法只收集真正的局部变量定义，不包括对外部作用域变量的修改

        Returns:
            本地符号的迭代器
        """
        # 使用 node_type 字符串判断，避免循环导入
        for stmt in self.statements:
            if stmt.node_type == "assign":
                yield from stmt.get_local_symbols()

            elif stmt.node_type in ("function", "method"):
                symbol = stmt.get_symbol()
                if symbol:
                    yield symbol

            # 处理嵌套语句块（不创建新作用域的语句块）
            elif stmt.node_type == "stmt_block":
                yield from stmt.get_local_symbols()

            # 处理控制流语句（if、for、while等不创建新作用域）
            elif not stmt.creates_scope():
                yield from stmt.get_local_symbols()

    def _extract_body_statements(self, body: List[ast.stmt]) -> List[IRNode]:
        """提取函数体语句，返回 IRNode 列表"""

        statements = []
        for stmt in body:
            if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant):
                # 跳过文档字符串
                continue

            ir_stmt = build_stmt_from_ast(stmt, self)
            if ir_stmt:
                statements.append(ir_stmt)

        return statements
